import cv2
import threading
import time
import os
import signal
import sys
import logging
import queue
from datetime import datetime

# 配置参数
DEFAULT_OUTPUT_DIR = "videos"
TARGET_FPS = 24
FRAME_INTERVAL = 1.0 / TARGET_FPS
QUEUE_MAX_SIZE = 100  # 帧队列最大容量

class VideoRecorder:
    def __init__(self, stream_url):
        self.stream_url = stream_url
        self.output_dir = DEFAULT_OUTPUT_DIR
        self.frame_queue = queue.Queue(maxsize=QUEUE_MAX_SIZE)
        self.is_recording = False
        self.is_streaming = False
        self.stream_thread = None
        self.record_thread = None
        self.last_frame = None  # 最后一帧用于填充
        self.start_time = 0.0
        self.frame_count = 0
        self.current_file = ""
        self.logger = self._setup_logger()
        self._stop_event = threading.Event()  # 线程同步事件

        # 创建输出目录
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            self.logger.info(f"创建目录: {self.output_dir}")

    def _setup_logger(self):
        """配置日志系统"""
        logger = logging.getLogger("VideoRecorder")
        logger.setLevel(logging.INFO)
        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
        
        # 控制台日志
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
        
        return logger

    def start_recording(self):
        """启动录制"""
        self.is_recording = True
        self.is_streaming = True
        self._stop_event.clear()
        self.start_time = time.time()
        self.frame_count = 0
        self.last_frame = None

        self.logger.info(f"开始录制视频流: {self.stream_url} (FPS={TARGET_FPS})")
        
        # 启动线程
        self.stream_thread = threading.Thread(target=self._process_stream)
        self.stream_thread.daemon = True
        self.stream_thread.start()

        self.record_thread = threading.Thread(target=self._record_frames)
        self.record_thread.daemon = True
        self.record_thread.start()

    def stop_recording(self):
        """安全停止录制并释放资源"""
        if not self.is_recording:
            return
            
        self.logger.info("收到停止信号，正在保存视频...")
        self.is_recording = False
        self.is_streaming = False
        self._stop_event.set()  # 通知线程停止
        
        # 清空队列
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except queue.Empty:
                pass
                
        # 等待线程结束
        if self.stream_thread and self.stream_thread.is_alive():
            self.stream_thread.join(timeout=2.0)
            
        if self.record_thread and self.record_thread.is_alive():
            self.record_thread.join(timeout=5.0)
            
        self.logger.info("录制已安全停止")

    def _process_stream(self):
        """处理视频流"""
        try:
            cap = cv2.VideoCapture(self.stream_url)
            if not cap.isOpened():
                self.logger.error("无法打开视频流")
                self.stop_recording()
                return
                
            self.logger.info("成功连接到视频流")
            
            while self.is_streaming and not self._stop_event.is_set():
                ret, frame = cap.read()
                if not ret:
                    self.logger.warning("读取帧失败，等待重连...")
                    time.sleep(1)
                    continue

                # 队列满时丢弃最旧帧，保留最新帧
                try:
                    self.frame_queue.put_nowait(frame)
                    self.last_frame = frame  # 更新最新帧
                except queue.Full:
                    try:
                        self.frame_queue.get_nowait()
                        self.frame_queue.put_nowait(frame)
                        self.last_frame = frame
                    except queue.Empty:
                        pass

                # 轻微休眠减少CPU占用
                time.sleep(0.001)
                
            cap.release()
            self.logger.info("视频流已关闭")
            
        except Exception as e:
            self.logger.error(f"流处理错误: {str(e)}")
        finally:
            self.is_streaming = False

    def _record_frames(self):
        """写入视频文件"""
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        out = None
        frame_size = None
        
        try:
            while self.is_recording and not self._stop_event.is_set():
                elapsed_time = time.time() - self.start_time
                expected_frames = int(elapsed_time * TARGET_FPS) + 1  # 确保帧数足够
                
                # 按固定帧率循环写入
                while self.frame_count < expected_frames and not self._stop_event.is_set():
                    if not self.is_recording:
                        break
                        
                    # 获取帧（优先使用队列中的新帧，否则用最后一帧填充）
                    if not self.frame_queue.empty():
                        frame = self.frame_queue.get()
                        self.last_frame = frame  # 更新最新帧
                    else:
                        if self.last_frame is None:
                            # 初始化前无帧，等待新帧
                            time.sleep(FRAME_INTERVAL)
                            continue
                        frame = self.last_frame  # 使用最后一帧填充

                    # 初始化写入器
                    if out is None:
                        height, width = frame.shape[:2]
                        frame_size = (width, height)
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        self.current_file = os.path.join(
                            self.output_dir,
                            f"recording_{timestamp}.mp4"
                        )
                        out = cv2.VideoWriter(
                            self.current_file, fourcc, TARGET_FPS, frame_size
                        )
                        self.logger.info(f"开始写入文件: {self.current_file}")

                    # 写入帧
                    out.write(frame)
                    self.frame_count += 1

                    # 控制帧率间隔
                    sleep_time = max(0, (self.frame_count / TARGET_FPS) - elapsed_time)
                    time.sleep(sleep_time)
                    
        except Exception as e:
            self.logger.error(f"录制错误: {str(e)}")
        finally:
            # 确保资源释放
            if out is not None:
                out.release()
                if os.path.exists(self.current_file):
                    file_size = os.path.getsize(self.current_file) / (1024 * 1024)
                    duration = time.time() - self.start_time
                    self.logger.info(f"文件保存成功: {self.current_file} ({file_size:.2f} MB, {duration:.2f}秒)")
                else:
                    self.logger.warning(f"文件未正确保存: {self.current_file}")
                    
            self.is_recording = False

def signal_handler(sig, frame):
    """Ctrl+C信号处理"""
    global recorder
    print("\n接收到中断信号，正在安全停止录制...")
    recorder.stop_recording()
    sys.exit(0)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"用法: python {sys.argv[0]} <视频流URL>")
        sys.exit(1)
        
    stream_url = sys.argv[1]
    recorder = VideoRecorder(stream_url)
    
    # 注册Ctrl+C信号处理
    signal.signal(signal.SIGINT, signal_handler)
    
    try:
        recorder.start_recording()
        # 主线程保持运行
        while True:
            time.sleep(0.1)
    except Exception as e:
        print(f"发生错误: {str(e)}")
        recorder.stop_recording()
        sys.exit(1)